#!/usr/bin/env python3
import argparse
import os
import sys
from time import sleep
from datetime import datetime
import grpc

# Import P4Runtime lib from parent utils dir
# Probably there's a better way of doing this.
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.switch import ShutdownAllSwitchConnections

SWITCH_TO_HOST_PORT = 1
S1_S2_PORT = 2
S1_S3_PORT = 3
S2_S1_PORT = 2
S2_S3_PORT  = 3
S3_S1_PORT = 2
S3_S2_PORT = 3


def writeTunnelRules(p4info_helper, ingress_sw, egress_sw, tunnel_id,
                     dst_eth_addr, dst_ip_addr, dst_port):
    """
    Installs three rules:
    1) An tunnel ingress rule on the ingress switch in the ipv4_lpm table that
       encapsulates traffic into a tunnel with the specified ID
    2) A transit rule on the ingress switch that forwards traffic based on
       the specified ID
    3) An tunnel egress rule on the egress switch that decapsulates traffic
       with the specified ID and sends it to the hostdst_port

    :param p4info_helper: the P4Info helper
    :param ingress_sw: the ingress switch connection
    :param egress_sw: the egress switch connection
    :param tunnel_id: the specified tunnel ID
    :param dst_eth_addr: the destination IP to match in the ingress rule
    :param dst_ip_addr: the destination Ethernet address to write in the
                        egress rule
    """
    # 1) Tunnel Ingress Rule
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={
            "hdr.ipv4.dstAddr": (dst_ip_addr, 32)
        },
        action_name="MyIngress.myTunnel_ingress",
        action_params={
            "dst_id": tunnel_id,
        })
    ingress_sw.WriteTableEntry(table_entry)
    print("Installed ingress tunnel rule on %s" % ingress_sw.name)

    # 2) Tunnel Transit Rule
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.myTunnel_exact",
        match_fields={
            "hdr.myTunnel.dst_id": tunnel_id
        },
        action_name="MyIngress.myTunnel_forward",
        action_params={
            "port": dst_port
        })
    ingress_sw.WriteTableEntry(table_entry)
    print("Installed transit tunnel rule on %s" % ingress_sw.name)

    # 3) Tunnel Egress Rule
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.myTunnel_exact",
        match_fields={
            "hdr.myTunnel.dst_id": tunnel_id
        },
        action_name="MyIngress.myTunnel_egress",
        action_params={
            "dstAddr": dst_eth_addr,
            "port": SWITCH_TO_HOST_PORT
        })
    egress_sw.WriteTableEntry(table_entry)
    print("Installed egress tunnel rule on %s" % egress_sw.name)


def readTableRules(p4info_helper, sw):
    """
    Reads the table entries from all tables on the switch.

    :param p4info_helper: the P4Info helper
    :param sw: the switch connection
    """
    print('\n----- Reading tables rules for %s -----' % sw.name)
    for response in sw.ReadTableEntries():
        for entity in response.entities:
            entry = entity.table_entry
            table_name = p4info_helper.get_tables_name(entry.table_id)
            print('%s: ' % table_name, end=' ')
            for m in entry.match:
                print(p4info_helper.get_match_field_name(table_name, m.field_id), end=' ')
                print('%r' % (p4info_helper.get_match_field_value(m),), end=' ')
            action = entry.action.action
            action_name = p4info_helper.get_actions_name(action.action_id)
            print('->', action_name, end=' ')
            for p in action.params:
                print(p4info_helper.get_action_param_name(action_name, p.param_id), end=' ')
                print('%r' % p.value, end=' ')
            print()

def logSwitchLinkCounter(p4info_helper, sw, counter_name, tunnel_id, ingress_sw_name, egress_sw_name, log_dir="link_logs"):
    """
    Logs counter data to a file named like 's1s2.txt', representing the switch-to-switch tunnel link.
    """
    os.makedirs(log_dir, exist_ok=True)
    # S1S2.txt...
    filename = f"{ingress_sw_name.upper()}{egress_sw_name.upper()}.txt"
    log_file = os.path.join(log_dir, filename)

    for response in sw.ReadCounters(p4info_helper.get_counters_id(counter_name), tunnel_id):
        for entity in response.entities:
            counter = entity.counter_entry
            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            line = f"[{timestamp}] {sw.name} {counter_name} {tunnel_id}: {counter.data.packet_count} packets ({counter.data.byte_count} bytes)\n"
            with open(log_file, "a") as f:
                f.write(line)

def printCounter(p4info_helper, sw, counter_name, index):
    """
    Reads the specified counter at the specified index from the switch. In our
    program, the index is the tunnel ID. If the index is 0, it will return all
    values from the counter.

    :param p4info_helper: the P4Info helper
    :param sw:  the switch connection
    :param counter_name: the name of the counter from the P4 program
    :param index: the counter index (in our case, the tunnel ID)
    """
    for response in sw.ReadCounters(p4info_helper.get_counters_id(counter_name), index):
        for entity in response.entities:
            counter = entity.counter_entry
            print("%s %s %d: %d packets (%d bytes)" % (
                sw.name, counter_name, index,
                counter.data.packet_count, counter.data.byte_count
            ))


def printGrpcError(e):
    print("gRPC Error:", e.details(), end=' ')
    status_code = e.code()
    print("(%s)" % status_code.name, end=' ')
    traceback = sys.exc_info()[2]
    print("[%s:%d]" % (traceback.tb_frame.f_code.co_filename, traceback.tb_lineno))


def main(p4info_file_path, bmv2_file_path):
    # Instantiate a P4Runtime helper from the p4info file
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        # Create a switch connection object for s1, s2, and s3
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1',
            address='127.0.0.1:50051',
            device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt')
        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2',
            address='127.0.0.1:50052',
            device_id=1,
            proto_dump_file='logs/s2-p4runtime-requests.txt')
        #add s3
        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3',
            address='127.0.0.1:50053',
            device_id=2,
            proto_dump_file='logs/s3-p4runtime-requests.txt')

        # Send master arbitration update message to establish this controller as master
        s1.MasterArbitrationUpdate()
        s2.MasterArbitrationUpdate()
        s3.MasterArbitrationUpdate()

        # Install the P4 program on the switches
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s1")
        s2.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s2")
        s3.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s3")

        # Write the rules that tunnel traffic from h1 to h2
        writeTunnelRules(p4info_helper, ingress_sw=s1, egress_sw=s2, tunnel_id=100,
                         dst_eth_addr="08:00:00:00:02:22", dst_ip_addr="10.0.2.2",dst_port = S1_S2_PORT)

        # Write the rules that tunnel traffic from h2 to h1
        writeTunnelRules(p4info_helper, ingress_sw=s2, egress_sw=s1, tunnel_id=101,
                         dst_eth_addr="08:00:00:00:01:11", dst_ip_addr="10.0.1.1",dst_port = S2_S1_PORT)

        # Write the rules that tunnel traffic from h1 to h3
        writeTunnelRules(p4info_helper, ingress_sw=s1, egress_sw=s3, tunnel_id=200,
                         dst_eth_addr="08:00:00:00:03:33", dst_ip_addr="10.0.3.3",dst_port = S1_S3_PORT)

        # Write the rules that tunnel traffic from h3 to h1
        writeTunnelRules(p4info_helper, ingress_sw=s3, egress_sw=s1, tunnel_id=201,
                         dst_eth_addr="08:00:00:00:01:11", dst_ip_addr="10.0.1.1",dst_port = S3_S1_PORT)

        # Write the rules that tunnel traffic from h2 to h3
        writeTunnelRules(p4info_helper, ingress_sw=s2, egress_sw=s3, tunnel_id=300,
                         dst_eth_addr="08:00:00:00:03:33", dst_ip_addr="10.0.3.3",dst_port = S2_S3_PORT)

        # Write the rules that tunnel traffic from h3 to h2
        writeTunnelRules(p4info_helper, ingress_sw=s3, egress_sw=s2, tunnel_id=301,
                         dst_eth_addr="08:00:00:00:02:22", dst_ip_addr="10.0.2.2",dst_port = S3_S2_PORT)

        # Read table entries from all switches
        readTableRules(p4info_helper, s1)
        readTableRules(p4info_helper, s2)
        readTableRules(p4info_helper, s3)

        # Print the tunnel counters every 2 seconds
        while True:
            sleep(2)
            print('\n----- Reading tunnel counters -----')

            # Tunnel 100: s1 (ingress) -> s2 (egress)  => S1S2.txt
            printCounter(p4info_helper, s1, "MyIngress.ingressTunnelCounter", 100)
            printCounter(p4info_helper, s2, "MyIngress.egressTunnelCounter", 100)
            logSwitchLinkCounter(p4info_helper, s1, "MyIngress.ingressTunnelCounter", 100, 's1', 's2')
            logSwitchLinkCounter(p4info_helper, s2, "MyIngress.egressTunnelCounter", 100, 's1', 's2')

            # Tunnel 101: s2 -> s1  => S2S1.txt
            printCounter(p4info_helper, s2, "MyIngress.ingressTunnelCounter", 101)
            printCounter(p4info_helper, s1, "MyIngress.egressTunnelCounter", 101)
            logSwitchLinkCounter(p4info_helper, s2, "MyIngress.ingressTunnelCounter", 101, 's2', 's1')
            logSwitchLinkCounter(p4info_helper, s1, "MyIngress.egressTunnelCounter", 101, 's2', 's1')

            # Tunnel 200: s1 -> s3  => S1S3.txt
            printCounter(p4info_helper, s1, "MyIngress.ingressTunnelCounter", 200)
            printCounter(p4info_helper, s3, "MyIngress.egressTunnelCounter", 200)
            logSwitchLinkCounter(p4info_helper, s1, "MyIngress.ingressTunnelCounter", 200, 's1', 's3')
            logSwitchLinkCounter(p4info_helper, s3, "MyIngress.egressTunnelCounter", 200, 's1', 's3')

            # Tunnel 201: s3 -> s1  => S3S1.txt
            printCounter(p4info_helper, s3, "MyIngress.ingressTunnelCounter", 201)
            printCounter(p4info_helper, s1, "MyIngress.egressTunnelCounter", 201)
            logSwitchLinkCounter(p4info_helper, s3, "MyIngress.ingressTunnelCounter", 201, 's3', 's1')
            logSwitchLinkCounter(p4info_helper, s1, "MyIngress.egressTunnelCounter", 201, 's3', 's1')

            # Tunnel 300: s2 -> s3  => S2S3.txt
            printCounter(p4info_helper, s2, "MyIngress.ingressTunnelCounter", 300)
            printCounter(p4info_helper, s3, "MyIngress.egressTunnelCounter", 300)
            logSwitchLinkCounter(p4info_helper, s2, "MyIngress.ingressTunnelCounter", 300, 's2', 's3')
            logSwitchLinkCounter(p4info_helper, s3, "MyIngress.egressTunnelCounter", 300, 's2', 's3')

            # Tunnel 301: s3 -> s2  => S3S2.txt
            printCounter(p4info_helper, s3, "MyIngress.ingressTunnelCounter", 301)
            printCounter(p4info_helper, s2, "MyIngress.egressTunnelCounter", 301)
            logSwitchLinkCounter(p4info_helper, s3, "MyIngress.ingressTunnelCounter", 301, 's3', 's2')
            logSwitchLinkCounter(p4info_helper, s2, "MyIngress.egressTunnelCounter", 301, 's3', 's2')




    except KeyboardInterrupt:
        print(" Shutting down.")
    except grpc.RpcError as e:
        printGrpcError(e)

    ShutdownAllSwitchConnections()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                        type=str, action="store", required=False,
                        default='./build/advanced_tunnel.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                        type=str, action="store", required=False,
                        default='./build/advanced_tunnel.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print("\np4info file not found: %s\nHave you run 'make'?" % args.p4info)
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print("\nBMv2 JSON file not found: %s\nHave you run 'make'?" % args.bmv2_json)
        parser.exit(1)
    main(args.p4info, args.bmv2_json)
